import torch
from torch import nn
from torch import Tensor
import torch.nn.functional as F
import numpy as np
from orbitP.script import config
from orbitP.layers.PatchTST_backbone import PatchTST_backbone
from orbitP.layers.PatchTST_layers import series_decomp
from orbitP.layers.FeatureEncoder import FeatureEncoder

class PatchTST(nn.Module):
    def __init__(self, args, **kwargs):

        super().__init__()
        max_seq_len = config.training_length
        d_k = None
        d_v = None
        norm = 'BatchNorm'
        attn_dropout = 0.
        act: str = "gelu"
        key_padding_mask:bool = 'auto'
        padding_var = None
        attn_mask = None
        res_attention: bool = True
        pre_norm: bool = False
        store_attn: bool = False
        pe: str = 'zeros'
        learn_pe: bool = True
        pretrain_head: bool = False
        head_type = 'flatten'
        verbose: bool = False


        # load parameters
        c_in = args.head_dim
        # c_in = args.feature_size+config.stampSize
        context_window = config.training_length
        target_window = config.predicting_length

        n_layers = args.num_layers
        n_heads = args.k
        d_model = args.d_model
        d_ff = args.d_ff
        dropout = args.dropout
        fc_dropout = args.dropout
        head_dropout = args.dropout

        individual = args.individual

        patch_len = args.patch_len
        stride = args.stride
        padding_patch = args.padding_patch

        revin = args.revin
        affine = args.affine
        subtract_last = args.subtract_last

        decomposition = args.decomposition
        kernel_size = args.kernel_size
        self.outputLayer = nn.Linear(c_in, config.outputSize+config.nll)
        self.FeatureEncoder = FeatureEncoder(args)

        # model
        self.decomposition = decomposition
        if self.decomposition:
            self.decomp_module = series_decomp(kernel_size)
            self.model_trend = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
                                                 patch_len=patch_len, stride=stride,
                                                 max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
                                                 n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
                                                 attn_dropout=attn_dropout,
                                                 dropout=dropout, act=act, key_padding_mask=key_padding_mask,
                                                 padding_var=padding_var,
                                                 attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
                                                 store_attn=store_attn,
                                                 pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout,
                                                 head_dropout=head_dropout, padding_patch=padding_patch,
                                                 pretrain_head=pretrain_head, head_type=head_type,
                                                 individual=individual, revin=revin, affine=affine,
                                                 subtract_last=subtract_last, verbose=verbose, **kwargs)
            self.model_res = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
                                               patch_len=patch_len, stride=stride,
                                               max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
                                               n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
                                               attn_dropout=attn_dropout,
                                               dropout=dropout, act=act, key_padding_mask=key_padding_mask,
                                               padding_var=padding_var,
                                               attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
                                               store_attn=store_attn,
                                               pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout,
                                               head_dropout=head_dropout, padding_patch=padding_patch,
                                               pretrain_head=pretrain_head, head_type=head_type, individual=individual,
                                               revin=revin, affine=affine,
                                               subtract_last=subtract_last, verbose=verbose, **kwargs)
        else:
            self.model = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
                                           patch_len=patch_len, stride=stride,
                                           max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
                                           n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
                                           attn_dropout=attn_dropout,
                                           dropout=dropout, act=act, key_padding_mask=key_padding_mask,
                                           padding_var=padding_var,
                                           attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
                                           store_attn=store_attn,
                                           pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout, head_dropout=head_dropout,
                                           padding_patch=padding_patch,
                                           pretrain_head=pretrain_head, head_type=head_type, individual=individual,
                                           revin=revin, affine=affine,
                                           subtract_last=subtract_last, verbose=verbose, **kwargs)

    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
        x_enc = x_enc.clone()
        x_mark_enc = x_mark_enc.clone()
        x_dec = x_dec[:,:,config.axis:-1].clone()
        x_mark_dec = x_mark_dec.clone()
        input = self.FeatureEncoder(x_enc, x_mark_enc, x_dec, x_mark_dec)
        if self.decomposition:
            res_init, trend_init = self.decomp_module(input)
            res_init, trend_init = res_init.permute(0, 2, 1), trend_init.permute(0, 2, 1)  # x: [Batch, Channel, Input length]
            res = self.model_res(res_init)
            trend = self.model_trend(trend_init)
            input = res + trend
            input = input.permute(0, 2, 1)  # x: [Batch, Input length, Channel]
        else:
            input = input.permute(0, 2, 1)  # x: [Batch, Channel, Input length]
            input = self.model(input)
            input = input.permute(0, 2, 1)  # x: [Batch, Input length, Channel]
        input = self.outputLayer(input)
        return input